"""
音频文件读写处理器
支持多种音频格式的读取和写入
"""

import io
import os
import tempfile
from typing import Optional, Tuple, Union, BinaryIO, Any
from pathlib import Path
import soundfile as sf
import numpy as np
from pydub import AudioSegment
import wave


class AudioFileHandler:
    """音频文件处理器"""
    
    def __init__(self):
        self.supported_formats = {
            'read': ['.wav', '.mp3', '.flac', '.ogg', '.aac', '.m4a', '.wma'],
            'write': ['.wav', '.mp3', '.flac', '.ogg']
        }
    
    def read_audio_file(
        self, 
        file_path: Union[str, Path, BinaryIO],
        sample_rate: Optional[int] = None,
        mono: bool = False,
        normalize: bool = False
    ) -> Tuple[np.ndarray, int]:
        """读取音频文件
        
        Args:
            file_path: 文件路径或文件对象
            sample_rate: 目标采样率，None表示保持原采样率
            mono: 是否转换为单声道
            normalize: 是否标准化音量
            
        Returns:
            Tuple[音频数据, 采样率]
        """
        try:
            if isinstance(file_path, (str, Path)):
                file_path = str(file_path)
                
                # 检查文件是否存在
                if not os.path.exists(file_path):
                    raise FileNotFoundError(f"Audio file not found: {file_path}")
                
                # 根据文件扩展名选择读取方法
                ext = os.path.splitext(file_path)[1].lower()
                
                if ext in ['.mp3', '.aac', '.m4a']:
                    return self._read_with_pydub(file_path, sample_rate, mono, normalize)
                else:
                    return self._read_with_soundfile(file_path, sample_rate, mono, normalize)
            
            else:
                # 处理文件对象
                return self._read_from_file_object(file_path, sample_rate, mono, normalize)
                
        except Exception as e:
            raise RuntimeError(f"Failed to read audio file: {str(e)}")
    
    def _read_with_soundfile(
        self, 
        file_path: str, 
        sample_rate: Optional[int],
        mono: bool,
        normalize: bool
    ) -> Tuple[np.ndarray, int]:
        """使用soundfile读取音频"""
        data, sr = sf.read(file_path)
        
        # 转换采样率
        if sample_rate is not None and sample_rate != sr:
            data = self._resample_audio(data, sr, sample_rate)
            sr = sample_rate
        
        # 转换为单声道
        if mono and len(data.shape) > 1:
            data = np.mean(data, axis=1)
        
        # 标准化
        if normalize:
            data = self._normalize_audio(data)
        
        return data, sr
    
    def _read_with_pydub(
        self, 
        file_path: str, 
        sample_rate: Optional[int],
        mono: bool,
        normalize: bool
    ) -> Tuple[np.ndarray, int]:
        """使用pydub读取音频"""
        audio = AudioSegment.from_file(file_path)
        
        # 转换采样率
        if sample_rate is not None:
            audio = audio.set_frame_rate(sample_rate)
            sr = sample_rate
        else:
            sr = audio.frame_rate
        
        # 转换为单声道
        if mono:
            audio = audio.set_channels(1)
        
        # 获取音频数据
        samples = np.array(audio.get_array_of_samples())
        
        # 处理立体声
        if audio.channels > 1:
            samples = samples.reshape((-1, audio.channels))
            samples = np.mean(samples, axis=1)
        
        # 转换为浮点数
        samples = samples.astype(np.float32) / np.iinfo(np.int16).max
        
        # 标准化
        if normalize:
            samples = self._normalize_audio(samples)
        
        return samples, sr
    
    def _read_from_file_object(
        self, 
        file_obj: BinaryIO,
        sample_rate: Optional[int],
        mono: bool,
        normalize: bool
    ) -> Tuple[np.ndarray, int]:
        """从文件对象读取音频"""
        # 保存到临时文件
        with tempfile.NamedTemporaryFile(delete=False) as temp_file:
            temp_file.write(file_obj.read())
            temp_file.seek(0)
            
            try:
                data, sr = self.read_audio_file(temp_file.name, sample_rate, mono, normalize)
                return data, sr
            finally:
                # 清理临时文件
                os.unlink(temp_file.name)
    
    def write_audio_file(
        self,
        data: np.ndarray,
        sample_rate: int,
        output_path: Union[str, Path],
        format: str = 'wav',
        bitrate: str = '192k',
        normalize: bool = False
    ) -> None:
        """写入音频文件
        
        Args:
            data: 音频数据
            sample_rate: 采样率
            output_path: 输出路径
            format: 音频格式
            bitrate: 比特率（仅对有损格式有效）
            normalize: 是否标准化音量
        """
        try:
            output_path = str(output_path)
            
            # 确保输出目录存在
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # 标准化
            if normalize:
                data = self._normalize_audio(data)
            
            # 根据格式选择写入方法
            if format.lower() in ['wav', 'flac', 'ogg']:
                self._write_with_soundfile(data, sample_rate, output_path, format)
            elif format.lower() in ['mp3']:
                self._write_with_pydub(data, sample_rate, output_path, format, bitrate)
            else:
                raise ValueError(f"Unsupported output format: {format}")
                
        except Exception as e:
            raise RuntimeError(f"Failed to write audio file: {str(e)}")
    
    def _write_with_soundfile(
        self,
        data: np.ndarray,
        sample_rate: int,
        output_path: str,
        format: str
    ):
        """使用soundfile写入音频"""
        sf.write(output_path, data, sample_rate, format=format)
    
    def _write_with_pydub(
        self,
        data: np.ndarray,
        sample_rate: int,
        output_path: str,
        format: str,
        bitrate: str
    ):
        """使用pydub写入音频"""
        # 转换为16位整数
        data_int16 = (data * np.iinfo(np.int16).max).astype(np.int16)
        
        # 创建AudioSegment
        audio = AudioSegment(
            data_int16.tobytes(),
            frame_rate=sample_rate,
            sample_width=data_int16.dtype.itemsize,
            channels=1 if len(data.shape) == 1 else data.shape[1]
        )
        
        # 导出文件
        audio.export(output_path, format=format, bitrate=bitrate)
    
    def get_audio_info(self, file_path: Union[str, Path, BinaryIO]) -> dict:
        """获取音频文件信息"""
        try:
            if isinstance(file_path, (str, Path)):
                file_path = str(file_path)
                
                if not os.path.exists(file_path):
                    raise FileNotFoundError(f"Audio file not found: {file_path}")
                
                # 获取基本信息
                info = sf.info(file_path)
                file_size = os.path.getsize(file_path)
                
                return {
                    'duration': info.duration,
                    'sample_rate': info.samplerate,
                    'channels': info.channels,
                    'format': info.format,
                    'subtype': info.subtype,
                    'frames': info.frames,
                    'file_size': file_size,
                    'bitrate': int(file_size * 8 / info.duration / 1000) if info.duration > 0 else 0
                }
            else:
                # 对于文件对象，需要先读取
                data, sr = self.read_audio_file(file_path)
                return {
                    'duration': len(data) / sr,
                    'sample_rate': sr,
                    'channels': 1 if len(data.shape) == 1 else data.shape[1],
                    'frames': len(data),
                    'format': 'unknown'
                }
                
        except Exception as e:
            raise RuntimeError(f"Failed to get audio info: {str(e)}")
    
    def _resample_audio(self, data: np.ndarray, orig_sr: int, target_sr: int) -> np.ndarray:
        """重采样音频"""
        try:
            import librosa
            return librosa.resample(data, orig_sr=orig_sr, target_sr=target_sr)
        except ImportError:
            # 如果没有librosa，使用简单的线性插值
            ratio = target_sr / orig_sr
            new_length = int(len(data) * ratio)
            indices = np.linspace(0, len(data) - 1, new_length)
            if len(data.shape) > 1:
                return np.array([np.interp(indices, np.arange(len(data)), data[:, ch]) 
                               for ch in range(data.shape[1])]).T
            else:
                return np.interp(indices, np.arange(len(data)), data)
    
    def _normalize_audio(self, data: np.ndarray) -> np.ndarray:
        """标准化音频音量"""
        if len(data) == 0:
            return data
        
        # 计算RMS
        rms = np.sqrt(np.mean(data ** 2))
        
        if rms > 0:
            # 标准化到-20dB
            target_rms = 0.1  # -20dB
            data = data * (target_rms / rms)
        
        # 限制幅度
        data = np.clip(data, -1.0, 1.0)
        
        return data
    
    def is_supported_format(self, file_path: Union[str, Path], operation: str = 'read') -> bool:
        """检查是否支持该格式"""
        if operation not in self.supported_formats:
            return False
        
        ext = os.path.splitext(str(file_path))[1].lower()
        return ext in self.supported_formats[operation]
    
    def convert_format(
        self,
        input_path: Union[str, Path, BinaryIO],
        output_path: Union[str, Path],
        target_format: str,
        sample_rate: Optional[int] = None,
        bitrate: str = '192k'
    ) -> None:
        """转换音频格式"""
        # 读取音频
        data, sr = self.read_audio_file(input_path, sample_rate=sample_rate)
        
        # 写入新格式
        self.write_audio_file(data, sr, output_path, target_format, bitrate)
    
    def extract_segment(
        self,
        file_path: Union[str, Path],
        start_time: float,
        end_time: float,
        sample_rate: Optional[int] = None
    ) -> Tuple[np.ndarray, int]:
        """提取音频片段"""
        data, sr = self.read_audio_file(file_path, sample_rate=sample_rate)
        
        start_sample = int(start_time * sr)
        end_sample = int(end_time * sr)
        
        # 确保索引在有效范围内
        start_sample = max(0, start_sample)
        end_sample = min(len(data), end_sample)
        
        return data[start_sample:end_sample], sr
    
    def concatenate_audio(self, audio_list: list) -> Tuple[np.ndarray, int]:
        """拼接多个音频"""
        if not audio_list:
            return np.array([]), 44100
        
        # 使用第一个音频的采样率
        sample_rate = audio_list[0][1]
        
        # 重采样所有音频到相同的采样率
        resampled_audios = []
        for audio, sr in audio_list:
            if sr != sample_rate:
                audio = self._resample_audio(audio, sr, sample_rate)
            resampled_audios.append(audio)
        
        # 拼接音频
        concatenated = np.concatenate(resampled_audios)
        
        return concatenated, sample_rate
